//Rijndael.cpp
#include "stdafx.h"
#include "Rijndael.h"


//Error Messages
char const* CRijndael::sm_szErrorMsg1 = "Object not Initialized";
char const* CRijndael::sm_szErrorMsg2 = "Data not multiple of Block Size";


//CONSTRUCTOR
CRijndael::CRijndael() : encryptKeys(NULL), decryptKeys(NULL), m_bKeyInit(false)
{
}

//DESTRUCTOR
CRijndael::~CRijndael()
{
	int i;
	if(encryptKeys != NULL)
	{
		for(i=0; i<=10; i++)
			delete [] encryptKeys[i];
		delete [] encryptKeys;
	}
	
	if(decryptKeys != NULL)
	{
		for(i=0; i<=10; i++)
			delete [] decryptKeys[i];
		delete [] decryptKeys;
	}
}

//Expand a user-supplied key material into a session key.
// key        - The 128 bit user-key to use.
// keylength  - 16 bytes
// blockSize  - The block size in bytes of this Rijndael (16 bytes).
void CRijndael::MakeKey(char const* key)
{
	int keylength = 16;
	int blockSize = 16;

	if(key == NULL)
		throw std::exception("Empty key");
	

	m_keylength = keylength; 
	m_blockSize = blockSize;
	

	int i;
	
	if(encryptKeys != NULL)
	{
		for(i=0; i<=10; i++)
			delete [] encryptKeys[i];
		delete [] encryptKeys;
	}
	
	if(decryptKeys != NULL)
	{
		for(i=0; i<=10; i++)
			delete [] decryptKeys[i];
		delete [] decryptKeys;
	}
	

	int BC = m_blockSize / 4;
	int j;

	encryptKeys = new int*[11]; //Encryption round keys
	for(i=0; i<=10; i++)
	{
		encryptKeys[i] = new int[BC];
		for(j=0; j<BC; j++)
			encryptKeys[i][j] = 0;
	}
	
	decryptKeys = new int*[11]; //Decryption round keys
	for(i=0; i<=10; i++)
	{
		decryptKeys[i] = new int[BC];
		for(j=0; j<BC; j++)
			decryptKeys[i][j] = 0;
	}
	
	int ROUND_KEY_COUNT = (11) * BC;
	int KC = m_keylength/4;
	int* tk = new int[KC];
	
	//Copy user material bytes into temporary ints
	int* pi = tk;
	char const* pc = key;
	for(i=0; i<KC; i++)
	{
		*pi = (unsigned char)*(pc++) << 24;
		*pi |= (unsigned char)*(pc++) << 16;
		*pi |= (unsigned char)*(pc++) << 8;
		*(pi++) |= (unsigned char)*(pc++);
	}


	//Copy values into round key arrays
	int t = 0;
	for(j=0; (j<KC)&&(t<ROUND_KEY_COUNT); j++,t++)
	{
		encryptKeys[t/BC][t%BC] = tk[j];
		decryptKeys[10 - (t/BC)][t%BC] = tk[j];
	}
	int tt, rconpointer = 0;
	while(t < ROUND_KEY_COUNT)
	{
		//Extrapolate using phi (the round key evolution function)
		tt = tk[KC-1];
		tk[0] ^= (SBox[(tt >> 16) & 0xFF] & 0xFF) << 24 ^
			(SBox[(tt >>  8) & 0xFF] & 0xFF) << 16 ^
			(SBox[ tt & 0xFF] & 0xFF) <<  8 ^
			(SBox[(tt >> 24) & 0xFF] & 0xFF) ^
			(sm_rcon[rconpointer++]  & 0xFF) << 24;
		if(KC != 8)
			for(i=1, j=0; i<KC;)
				tk[i++] ^= tk[j++];
		else
		{
			for(i=1, j=0; i<KC/2; )
				tk[i++] ^= tk[j++];
			tt = tk[KC/2-1];
			tk[KC/2] ^= (SBox[ tt & 0xFF] & 0xFF) ^
				(SBox[(tt >>  8) & 0xFF] & 0xFF) <<  8 ^
				(SBox[(tt >> 16) & 0xFF] & 0xFF) << 16 ^
				(SBox[(tt >> 24) & 0xFF] & 0xFF) << 24;
			for(j = KC/2, i=j+1; i<KC; )
				tk[i++] ^= tk[j++];
		}
	
		//Copy values into round key arrays
		for(j=0; (j<KC) && (t<ROUND_KEY_COUNT); j++, t++)
		{
			encryptKeys[t/BC][t%BC] = tk[j];
			decryptKeys[10 - (t/BC)][t%BC] = tk[j];
		}
	}

	delete [] tk;
	
	//Inverse MixColumn where needed
	for(int r=1; r<10; r++)
		for(j=0; j<BC; j++)
		{
			tt = decryptKeys[r][j];
			decryptKeys[r][j] = sm_U1[(tt >> 24) & 0xFF] ^
				sm_U2[(tt >> 16) & 0xFF] ^
				sm_U3[(tt >>  8) & 0xFF] ^
				sm_U4[tt & 0xFF];
		}
	m_bKeyInit = true;
}

//Convenience method to encrypt exactly one block of plaintext, assuming
//Rijndael's default block size (128-bit).
// in         - The plaintext
// result     - The ciphertext generated from a plaintext using the key
void CRijndael::DefEncryptBlock(char const* in, char* result)
{
	if(false==m_bKeyInit)
		throw std::exception(sm_szErrorMsg1);
	int* Ker = encryptKeys[0];
	//cout<<"Entered DefEncryptBlock \n";

	int t0 = ((unsigned char)*(in++) << 24);
	t0 |= ((unsigned char)*(in++) << 16);
	t0 |= ((unsigned char)*(in++) << 8);
	(t0 |= (unsigned char)*(in++)) ^= Ker[0];
	int t1 = ((unsigned char)*(in++) << 24);
	t1 |= ((unsigned char)*(in++) << 16);
	t1 |= ((unsigned char)*(in++) << 8);
	(t1 |= (unsigned char)*(in++)) ^= Ker[1];
	int t2 = ((unsigned char)*(in++) << 24);
	t2 |= ((unsigned char)*(in++) << 16);
	t2 |= ((unsigned char)*(in++) << 8);
	(t2 |= (unsigned char)*(in++)) ^= Ker[2];
	int t3 = ((unsigned char)*(in++) << 24);
	t3 |= ((unsigned char)*(in++) << 16);
	t3 |= ((unsigned char)*(in++) << 8);
	(t3 |= (unsigned char)*(in++)) ^= Ker[3];
	int a0, a1, a2, a3;
	//Apply Round Transforms
	for (int r = 1; r < 10; r++)
	{
		Ker = encryptKeys[r];
		a0 = (sm_T1[(t0 >> 24) & 0xFF] ^
			sm_T2[(t1 >> 16) & 0xFF] ^
			sm_T3[(t2 >>  8) & 0xFF] ^
			sm_T4[t3 & 0xFF]) ^ Ker[0];
		a1 = (sm_T1[(t1 >> 24) & 0xFF] ^
			sm_T2[(t2 >> 16) & 0xFF] ^
			sm_T3[(t3 >>  8) & 0xFF] ^
			sm_T4[t0 & 0xFF]) ^ Ker[1];
		a2 = (sm_T1[(t2 >> 24) & 0xFF] ^
			sm_T2[(t3 >> 16) & 0xFF] ^
			sm_T3[(t0 >>  8) & 0xFF] ^
			sm_T4[t1 & 0xFF]) ^ Ker[2];
		a3 = (sm_T1[(t3 >> 24) & 0xFF] ^
			sm_T2[(t0 >> 16) & 0xFF] ^
			sm_T3[(t1 >>  8) & 0xFF] ^
			sm_T4[t2 & 0xFF]) ^ Ker[3];
		t0 = a0;
		t1 = a1;
		t2 = a2;
		t3 = a3;
	}
	//Last Round is special
	Ker = encryptKeys[10];
	int tt = Ker[0];
	result[0] = SBox[(t0 >> 24) & 0xFF] ^ (tt >> 24);
	result[1] = SBox[(t1 >> 16) & 0xFF] ^ (tt >> 16);
	result[2] = SBox[(t2 >>  8) & 0xFF] ^ (tt >>  8);
	result[3] = SBox[t3 & 0xFF] ^ tt;
	tt = Ker[1];
	result[4] = SBox[(t1 >> 24) & 0xFF] ^ (tt >> 24);
	result[5] = SBox[(t2 >> 16) & 0xFF] ^ (tt >> 16);
	result[6] = SBox[(t3 >>  8) & 0xFF] ^ (tt >>  8);
	result[7] = SBox[t0 & 0xFF] ^ tt;
	tt = Ker[2];
	result[8] = SBox[(t2 >> 24) & 0xFF] ^ (tt >> 24);
	result[9] = SBox[(t3 >> 16) & 0xFF] ^ (tt >> 16);
	result[10] = SBox[(t0 >>  8) & 0xFF] ^ (tt >>  8);
	result[11] = SBox[t1 & 0xFF] ^ tt;
	tt = Ker[3];
	result[12] = SBox[(t3 >> 24) & 0xFF] ^ (tt >> 24);
	result[13] = SBox[(t0 >> 16) & 0xFF] ^ (tt >> 16);
	result[14] = SBox[(t1 >>  8) & 0xFF] ^ (tt >>  8);
	result[15] = SBox[t2 & 0xFF] ^ tt;

	/*for(int i=0; i<=15; i++)
	{
		cout<<"The result["<<i<<"] ="<<(int)result[i]<<endl;
	}*/

}

//Convenience method to decrypt exactly one block of plaintext, assuming
//Rijndael's default block size (128-bit).
// in         - The ciphertext.
// result     - The plaintext generated from a ciphertext using the session key.
void CRijndael::DefDecryptBlock(char const* in, char* result)
{
	if(false==m_bKeyInit)
		throw std::exception(sm_szErrorMsg1);
	int* Kdr = decryptKeys[0];
	int t0 = ((unsigned char)*(in++) << 24);
	t0 = t0 | ((unsigned char)*(in++) << 16);
	t0 |= ((unsigned char)*(in++) << 8);
	(t0 |= (unsigned char)*(in++)) ^= Kdr[0];
	int t1 = ((unsigned char)*(in++) << 24);
	t1 |= ((unsigned char)*(in++) << 16);
	t1 |= ((unsigned char)*(in++) << 8);
	(t1 |= (unsigned char)*(in++)) ^= Kdr[1];
	int t2 = ((unsigned char)*(in++) << 24);
	t2 |= ((unsigned char)*(in++) << 16);
	t2 |= ((unsigned char)*(in++) << 8);
	(t2 |= (unsigned char)*(in++)) ^= Kdr[2];
	int t3 = ((unsigned char)*(in++) << 24);
	t3 |= ((unsigned char)*(in++) << 16);
	t3 |= ((unsigned char)*(in++) << 8);
	(t3 |= (unsigned char)*(in++)) ^= Kdr[3];
	int a0, a1, a2, a3;
	for(int r = 1; r < 10 ; r++) // apply round transforms
	{
		Kdr = decryptKeys[r];
		a0 = (sm_T5[(t0 >> 24) & 0xFF] ^
			sm_T6[(t3 >> 16) & 0xFF] ^
			sm_T7[(t2 >>  8) & 0xFF] ^
			sm_T8[ t1        & 0xFF] ) ^ Kdr[0];
		a1 = (sm_T5[(t1 >> 24) & 0xFF] ^
			sm_T6[(t0 >> 16) & 0xFF] ^
			sm_T7[(t3 >>  8) & 0xFF] ^
			sm_T8[ t2        & 0xFF] ) ^ Kdr[1];
		a2 = (sm_T5[(t2 >> 24) & 0xFF] ^
			sm_T6[(t1 >> 16) & 0xFF] ^
			sm_T7[(t0 >>  8) & 0xFF] ^
			sm_T8[ t3        & 0xFF] ) ^ Kdr[2];
		a3 = (sm_T5[(t3 >> 24) & 0xFF] ^
			sm_T6[(t2 >> 16) & 0xFF] ^
			sm_T7[(t1 >>  8) & 0xFF] ^
			sm_T8[ t0        & 0xFF] ) ^ Kdr[3];
		t0 = a0;
		t1 = a1;
		t2 = a2;
		t3 = a3;
	}
	//Last Round is special
	Kdr = decryptKeys[10];
	int tt = Kdr[0];
	result[ 0] = SBoxInv[(t0 >> 24) & 0xFF] ^ (tt >> 24);
	result[ 1] = SBoxInv[(t3 >> 16) & 0xFF] ^ (tt >> 16);
	result[ 2] = SBoxInv[(t2 >>  8) & 0xFF] ^ (tt >>  8);
	result[ 3] = SBoxInv[ t1 & 0xFF] ^ tt;
	tt = Kdr[1];
	result[ 4] = SBoxInv[(t1 >> 24) & 0xFF] ^ (tt >> 24);
	result[ 5] = SBoxInv[(t0 >> 16) & 0xFF] ^ (tt >> 16);
	result[ 6] = SBoxInv[(t3 >>  8) & 0xFF] ^ (tt >>  8);
	result[ 7] = SBoxInv[ t2 & 0xFF] ^ tt;
	tt = Kdr[2];
	result[ 8] = SBoxInv[(t2 >> 24) & 0xFF] ^ (tt >> 24);
	result[ 9] = SBoxInv[(t1 >> 16) & 0xFF] ^ (tt >> 16);
	result[10] = SBoxInv[(t0 >>  8) & 0xFF] ^ (tt >>  8);
	result[11] = SBoxInv[ t3 & 0xFF] ^ tt;
	tt = Kdr[3];
	result[12] = SBoxInv[(t3 >> 24) & 0xFF] ^ (tt >> 24);
	result[13] = SBoxInv[(t2 >> 16) & 0xFF] ^ (tt >> 16);
	result[14] = SBoxInv[(t1 >>  8) & 0xFF] ^ (tt >>  8);
	result[15] = SBoxInv[ t0 & 0xFF] ^ tt;
}

//Encrypt exactly one block of plaintext.
// in           - The plaintext.
// result       - The ciphertext generated from a plaintext using the key.
void CRijndael::EncryptBlock(char const* in, char* result)
{
	if(false==m_bKeyInit)
		throw std::exception(sm_szErrorMsg1);
    if(DEFAULT_BLOCK_SIZE == m_blockSize)
	{
		std::cout<<"default block size :"<< m_blockSize;
		DefEncryptBlock(in, result);
		return;
	}
	int BC = m_blockSize / 4;
	int SC = (BC == 4) ? 0 : (BC == 6 ? 1 : 2);
	int s1 = sm_shifts[SC][1][0];
	int s2 = sm_shifts[SC][2][0];
	int s3 = sm_shifts[SC][3][0];
	//Temporary Work Arrays
	int* a = new int[BC];
	int* t = new int[BC];
	int i;
	int tt;
	int* pi = t;
	for(i=0; i<BC; i++)
	{
		*pi = ((unsigned char)*(in++) << 24);
		*pi |= ((unsigned char)*(in++) << 16);
		*pi |= ((unsigned char)*(in++) << 8);
		(*(pi++) |= (unsigned char)*(in++)) ^= encryptKeys[0][i];
	}
	//Apply Round Transforms
	for(int r=1; r<10; r++)
	{
		for(i=0; i<BC; i++)
			a[i] = (sm_T1[(t[i] >> 24) & 0xFF] ^
				sm_T2[(t[(i + s1) % BC] >> 16) & 0xFF] ^
				sm_T3[(t[(i + s2) % BC] >>  8) & 0xFF] ^
				sm_T4[ t[(i + s3) % BC] & 0xFF] ) ^ encryptKeys[r][i];
		memcpy(t, a, 4*BC);
	}
	int j;
	//Last Round is Special
	for(i=0,j=0; i<BC; i++)
	{
		tt = encryptKeys[10][i];
		//cout<<"The result["<<j<<"] ="<<result[j];
		result[j++] = SBox[(t[i] >> 24) & 0xFF] ^ (tt >> 24);
		//cout<<"The result["<<j<<"] ="<<result[j];
		result[j++] = SBox[(t[(i + s1) % BC] >> 16) & 0xFF] ^ (tt >> 16);
		//cout<<"The result["<<j<<"] ="<<result[j];
		result[j++] = SBox[(t[(i + s2) % BC] >>  8) & 0xFF] ^ (tt >>  8);
		//cout<<"The result["<<j<<"] ="<<result[j];
		result[j++] = SBox[ t[(i + s3) % BC] & 0xFF] ^ tt;
		//cout<<"The result["<<j<<"] ="<<result[j];
	}
	delete [] a;
	delete [] t;
}

//Decrypt exactly one block of ciphertext.
// in         - The ciphertext.
// result     - The plaintext generated from a ciphertext using the session key.
void CRijndael::DecryptBlock(char const* in, char* result)
{
	if(false==m_bKeyInit)
		throw std::exception(sm_szErrorMsg1);
	if(DEFAULT_BLOCK_SIZE == m_blockSize)
	{
		DefDecryptBlock(in, result);
		return;
	}
	int BC = m_blockSize / 4;
	int SC = BC == 4 ? 0 : (BC == 6 ? 1 : 2);
	int s1 = sm_shifts[SC][1][1];
	int s2 = sm_shifts[SC][2][1];
	int s3 = sm_shifts[SC][3][1];
	//Temporary Work Arrays
	int* a = new int[BC];
	int* t = new int[BC];
	int i;
	int tt;
	int* pi = t;
	for(i=0; i<BC; i++)
	{
		*pi = ((unsigned char)*(in++) << 24);
		*pi |= ((unsigned char)*(in++) << 16);
		*pi |= ((unsigned char)*(in++) << 8);
		(*(pi++) |= (unsigned char)*(in++)) ^= decryptKeys[0][i];
	}
	//Apply Round Transforms
	for(int r=1; r<10; r++)
	{
		for(i=0; i<BC; i++)
			a[i] = (sm_T5[(t[i] >> 24) & 0xFF] ^
				sm_T6[(t[(i + s1) % BC] >> 16) & 0xFF] ^
				sm_T7[(t[(i + s2) % BC] >>  8) & 0xFF] ^
				sm_T8[ t[(i + s3) % BC] & 0xFF]) ^ decryptKeys[r][i];
		memcpy(t, a, 4*BC);
	}
	int j;
	//Last Round is Special
	for(i=0,j=0; i<BC; i++)
	{
		tt = decryptKeys[10][i];
		result[j++] = SBoxInv[(t[i] >> 24) & 0xFF] ^ (tt >> 24);
		result[j++] = SBoxInv[(t[(i + s1) % BC] >> 16) & 0xFF] ^ (tt >> 16);
		result[j++] = SBoxInv[(t[(i + s2) % BC] >>  8) & 0xFF] ^ (tt >>  8);
		result[j++] = SBoxInv[ t[(i + s3) % BC] & 0xFF] ^ tt;
	}
	delete [] a;
	delete [] t;
}

void CRijndael::EncryptString(char const* in, char* result, size_t n)
{
	memset(result,0,300);
	if(false==m_bKeyInit)
		throw std::exception(sm_szErrorMsg1);
	//n should be > 0 and multiple of m_blockSize
	if(0==n || n%m_blockSize!=0)
		throw std::exception(sm_szErrorMsg2);
	unsigned int i;
	char const* pin;
	char* presult;
	for(i=0,pin=in,presult=result; i<(n/m_blockSize); i++)
	{
		EncryptBlock(pin, presult);
		pin += m_blockSize;
		presult += m_blockSize;
	}

	result[n] = '\0';
	
}

void CRijndael::DecryptString(char const* in, char* result, size_t n)
{
	if(false==m_bKeyInit)
		throw std::exception(sm_szErrorMsg1);
	//n should be > 0 and multiple of m_blockSize
	//if(0==n || n%m_blockSize!=0)
	//	throw exception(sm_szErrorMsg2);
	unsigned int i;
	char const* pin;
	char* presult;
	
	for(i=0,pin=in,presult=result; i<n/m_blockSize; i++)
	{
		DecryptBlock(pin, presult);
		pin += m_blockSize;
		presult += m_blockSize;
	}
	
	result[n] = '\0';
}

